from collections import defaultdict
import json
import re
# import openai
from model import deepseek, qwen_72B, chatgpt
import concurrent.futures
import threading

# # OpenAI全局配置
# openai.base_url = "https://free.v36.cm/v1/"
# openai.default_headers = {"x-foo": "true"}
# # openai.api_key = "sk-ygs1G1n4rb74v8fJ51F6B4D3749a4d978aB99b4308D2A993"
# openai.api_key = "sk-GbQ7V9AKYPUSLsu88e76A5C7B152449cAbFa0f5eCbD18aE3"

def read(file_path):
    data = defaultdict(dict)
    with open(file_path, 'r', encoding='utf-8') as file:
        lines = file.readlines()
    
    for index, line in enumerate(lines):
        tmp = json.loads(line)
        data[tmp['key']] = tmp
    return data

def process(data, ori_tag, tgt_tag):
    lock = threading.Lock()
    
    def process_function(key, txt):
        try:
            prompt = f'原文​​："{txt}"\n任务​​：这是对某一道数学题目的分析与回复。现在我希望你能提取这段文字中的最终答案。你只需给出结果。以下是几个例子：\n给定文本："如果文具盒每个22元，那么买3个文具盒的总价就是 (22 * 3 = 66) 元。所以，买3个文具盒一共需要66元。"，这段文字中最终答案是"66"，则你只需要回复"66"。\n给定文本："甲乙两队一共完成了总工程的 \( \\frac{{2}}{{3}} \)。"，这段文字中最终答案是"\\frac{{2}}{{3}}"，则你只需要回复"\\frac{{2}}{{3}}"。\n给定文本："the quantity eight minus ninety divided by ninety five equals negative one point零五二六三."，这段文字中最终答案是"negative one point零五二六三"，则你只需要回复"negative one point零五二六三"。\n注意，你无需其他文字的内容，且你的回复必须在原文中出现过！！即可以通过(answer in txt)判断。'
            # n_txt = deepseek(prompt)
            idx = 0
            # n_txt = None
            while idx < 3:
                try:
                    while True:
                        try:
                            if idx == 0:
                                n_txt = qwen_72B(prompt)
                            elif idx == 1:
                                n_txt = chatgpt(prompt)
                            elif idx == 2:
                                n_txt = deepseek(prompt)
                            break
                        except:
                            pass
                    print(rf'[{key}] [{txt}] -> [{n_txt}]')
                    assert (n_txt in txt) or (re.sub(r'\\boxed{(.*)}', r'\1', n_txt, flags=re.DOTALL) in txt) or (re.sub(r'\boxed{(.*)}', r'\1', n_txt, flags=re.DOTALL) in txt)
                    break
                except:
                    idx += 1
            assert idx != 3
            with lock:
                data[key][tgt_tag] = n_txt
        except Exception as e:
            print(f'{key} error')
    
    with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
        futures = []
        for key, value in data.items():
            txt = value[ori_tag]
            futures.append(executor.submit(process_function, key, txt))
        
        for future in concurrent.futures.as_completed(futures):
            pass  # 异常已在process函数中处理

    return data

def write(file_path, error_path, data, tgt_tag):
    with open(file_path, 'w', encoding='utf-8') as file, open(error_path, 'w', encoding='utf-8') as error:
        for key, value in data.items():
            if tgt_tag in value.keys():
                json.dump(value, file, ensure_ascii=False)
                file.write('\n')
            else:
                # json.dump(value, file, ensure_ascii=False)
                value[tgt_tag] = 'ERROR'
                json.dump(value, error, ensure_ascii=False)
                error.write('\n')

if __name__ == '__main__':
    file_path = '/mnt/sfs/asr/code/osum_xlgeng_3B/voicebench_wsy/qwen2_omni_chat/1.jsonl'
    out_path = file_path + '.n'
    err_path = file_path + '.e'
    ori_tag = 'predict'
    tgt_tag = 'predict_answer'

    data = read(file_path)
    data_n = process(data, ori_tag, tgt_tag)
    write(out_path, err_path, data_n, tgt_tag)